#!/usr/bin/python
# coding=utf-8

import sys
import psycopg2
import time
import re
import sqlparse
from sqlparse.sql import Identifier, Where, Comparison, Parenthesis, Function

# 连接即将执行query的数据库
def conn_db():
    # print '[{0}] 连接数据库 ...'.format(time.strftime('%X'))
    conn_str1 = 'host=192.168.236.191 dbname=postgres user=postgres'
    conn = psycopg2.connect(conn_str1)
    cur = conn.cursor()
    return (cur)


# 对query语句进行简单的分类,并对query语句分类计算总数
def analyze_query(sql_text):
    sql_text = sql_text.replace('\n', '')
    sql_text = sql_text.replace('\r', '')
    sql_list = sqlparse.parse(sql_text)
    sum_u = 0
    sum_i = 0
    sum_d = 0
    sum_s = 0
    sum_o = 0
    sum_a = len(sql_list)
    report_list = ''
    report = ''
    for sql in sql_list:
        sql_type = sql.get_type()
        print(sql_type)
        if sql_type == "INSERT":
            sum_i = sum_i + 1
            report = insert_query(sql)
        elif sql_type == "UPDATE":
            sum_u = sum_u + 1
            report = update_query(sql)
        elif sql_type == "DELETE":
            sum_d = sum_d + 1
            report = delete_query(sql)
        elif sql_type == "SELECT":
            sum_s = sum_s + 1
        else:
            sum_o += 1
            print(sql)
        report_list = report_list + '\n' + report
    if report_list.strip(' ').strip('\n') is not '':
        checkreport = ' insert语句 :%s条 \n update语句 :%s条 \n delete语句 :%s条 \n select语句 :%s条 \n 非DML语句  :%s条 \n report: %s' % (
        sum_i, sum_u, sum_d, sum_s, sum_o, report_list)
        print checkreport
    else:
        checkreport = ' insert语句 :%s条 \n update语句 :%s条 \n delete语句 :%s条 \n select语句 :%s条 \n 非DML语句  :%s条 \n report: success' % (
        sum_i, sum_u, sum_d, sum_s, sum_o)
        print checkreport


# INSERT语句处理规则：只判断tablename、into、values关键字是否而存在
def insert_query(sql):
    tablename = ''
    into_key_word = ''
    values_key_work = ''
    report = ''
    for token in sql.tokens:
        if isinstance(token, Identifier) or isinstance(token, Function):
            tablename = str(token).strip(' ').split('(')[0]
        if token.is_keyword:
            if str(token) == 'into':
                into_key_word = token
            if str(token) == 'values':
                values_key_work = token
    if (tablename is '') or (into_key_word is '') or (values_key_work is ''):
        report = '关键字 insert or into or values 不存在:%s' % str(sql)
    else:
        report = ''
    return report


# UPDATE语句处理规则：需要判断表和主键的对应关系是否正确，set关键字是否存在
def update_query(sql):
    tablename = ''
    primary_key = ''
    set_key_word = ''
    cur = conn_db()
    for token in sql.tokens:
        print type(token)
        if isinstance(token, Identifier):
            tablename = str(token)
            print(tablename)
        if token.is_keyword:
            if str(token) == 'set':
                set_key_word = str(token)
        if isinstance(token, Where):
            primary_key = str(token).split(' ')[1].split('=')[0]
            print(primary_key)
    print "table=%s, set_key_word=%s, primary_key=%s" % (tablename, set_key_word, primary_key)
    if (primary_key is not '') and (tablename is not '') and (set_key_word is not ''):
        sql = "select primary_key from primary_relation where tablename='%s'" % tablename
        try:
            cur.execute(sql)
            value, = cur.fetchall()
            primary_key_tmp = value[0]
            if str(primary_key_tmp) == str(primary_key):
                report = ''
            else:
                report = '表名与主键不匹配:%s' % str(sql)
        except:
            report = '表名与主键不匹配:%s' % str(sql)
    elif (primary_key is not '') and (tablename is not ''):
        report = 'SET 关键字不存在     :%s' % str(sql)
    elif (primary_key is not '') and (set_key_word is not ''):
        report = '表名不存在          :%s' % str(sql)
    elif (tablename is not '') and (set_key_word is not ''):
        report = '主键字段或WHERE不存在:%s' % str(sql)
    else:
        report = '这是SQL吗            :%s' % str(sql)


    return report

# DELETE语句处理规则：判断表与主键的对应关系是否正确，from关键字是否存在
def delete_query(sql):
    tablename = ''
    primary_key = ''
    delete_key_word = ''
    from_key_word = ''
    cur = conn_db()
    for token in sql.tokens:
        if isinstance(token, Identifier) or isinstance(token, Function):
            tablename = str(token).strip(' ').split('(')[0]
        if token.is_keyword:
            if str(token) == 'delete':
                delete_key_word = token
            if str(token) == 'from':
                from_key_word = token
        if isinstance(token, Where):
            primary_key = str(token).split(' ')[1].split('=')[0]

    if (primary_key is not '') and (tablename is not '') and (from_key_word is not '') and (delete_key_word is not ''):
        sql = "select primary_key from primary_relation where tablename='%s'" % tablename
        cur.execute(sql)
        value, = cur.fetchall()
        primary_key_tmp = value[0]
        if str(primary_key_tmp) == str(primary_key):
            report = ''
        else:
            report = '主键不匹配或模式名不存在:%s' % str(sql)
    else:
        report = '关键字不存在         :%s' % str(sql)
    return report

if __name__ == "__main__":
    analyze_query(sys.argv[1])

